# LangGraph 多工具智能体：天气查询 + 计算器
# 功能：根据用户输入智能判断并调用对应工具

from typing import Annotated, Literal, Optional
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain.chat_models import init_chat_model
from langchain.schema import AIMessage, SystemMessage
from langchain.tools import tool
import json
import re


# ========== 1. 定义工具 ==========
@tool
def get_weather(city: str) -> str:
    """
    查询指定城市的天气信息
    参数：
        city: 城市名称（如：北京、上海、深圳）
    返回：
        天气信息字符串
    """
    # 模拟天气数据（实际项目中可调用真实天气 API，如高德、和风天气等）
    weather_data = {
        "北京": "晴天，气温 18-28℃，东南风 2-3级，空气质量良",
        "上海": "多云转晴，气温 22-30℃，东风 3-4级，空气质量优",
        "深圳": "雷阵雨，气温 26-32℃，南风 2级，空气质量良",
        "广州": "阴转小雨，气温 25-31℃，无持续风向，空气质量良",
        "杭州": "晴，气温 20-29℃，东北风 1-2级，空气质量优"
    }

    result = weather_data.get(city, f"暂无{city}的天气数据，仅支持：北京、上海、深圳、广州、杭州")
    return f"【天气查询】{city}：{result}"


@tool
def calculator(expression: str) -> str:
    """
    计算数学表达式
    参数：
        expression: 数学表达式（如：2+3、10*5、100/4）
    返回：
        计算结果
    """
    try:
        # 安全评估数学表达式（仅允许数字和基本运算符）
        # 移除空格并验证表达式安全性
        clean_expr = expression.replace(" ", "")
        if not re.match(r'^[\d+\-*/().]+$', clean_expr):
            return "【计算器】错误：表达式包含非法字符，仅支持数字和 +-*/() 运算符"

        result = eval(clean_expr)
        return f"【计算器】{expression} = {result}"
    except Exception as e:
        return f"【计算器】计算错误：{str(e)}"


# ========== 2. 定义 State ==========
class ToolAgentState(TypedDict):
    messages: Annotated[list, add_messages]  # 对话历史
    tool_needed: Optional[str]  # 需要的工具类型："weather" | "calculator" | "none"
    tool_input: Optional[str]  # 工具输入参数
    tool_result: Optional[str]  # 工具执行结果


# ========== 3. 初始化 LLM ==========
llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")


# ========== 4. 定义节点 ==========

# 节点1：意图识别 - 判断需要什么工具
def intent_recognition_node(state: ToolAgentState) -> dict:
    """
    分析用户输入，判断需要调用哪个工具
    返回：tool_needed（"weather" | "calculator" | "none"）和 tool_input
    """
    # 获取最新的用户消息
    user_message = next(
        msg for msg in reversed(state["messages"])
        if hasattr(msg, 'type') and msg.type == "human"
    ).content

    # 构造意图识别提示词
    intent_prompt = SystemMessage(
        content=f"""你是一个工具调用判断器，分析用户问题并返回 JSON 格式结果。

规则：
1. 如果问题是查询天气（如"北京天气""上海明天下雨吗"），返回：
   {{"tool_needed": "weather", "tool_input": "城市名"}}

2. 如果问题是数学计算（如"3+5等于多少""计算100除以4"），返回：
   {{"tool_needed": "calculator", "tool_input": "数学表达式"}}

3. 如果是其他问题（聊天、常识等），返回：
   {{"tool_needed": "none", "tool_input": ""}}

用户问题：{user_message}

仅返回 JSON，不要附加其他内容！"""
    )

    # 调用 LLM 进行意图识别
    response = llm.invoke([intent_prompt])

    # 解析结果
    try:
        result = json.loads(response.content)
        tool_needed = result.get("tool_needed", "none")
        tool_input = result.get("tool_input", "")
    except:
        # 解析失败，默认不调用工具
        tool_needed = "none"
        tool_input = ""

    print(f"[意图识别] 工具: {tool_needed}, 参数: {tool_input}")

    return {
        "tool_needed": tool_needed,
        "tool_input": tool_input
    }


# 节点2：天气工具节点
def weather_tool_node(state: ToolAgentState) -> dict:
    """调用天气查询工具"""
    city = state["tool_input"]
    print(f"[工具调用] 正在查询 {city} 的天气...")

    result = get_weather.invoke(city)

    return {"tool_result": result}


# 节点3：计算器工具节点
def calculator_tool_node(state: ToolAgentState) -> dict:
    """调用计算器工具"""
    expression = state["tool_input"]
    print(f"[工具调用] 正在计算 {expression}...")

    result = calculator.invoke(expression)

    return {"tool_result": result}


# 节点4：生成回答节点
def generate_answer_node(state: ToolAgentState) -> dict:
    """
    基于对话历史和工具结果（如有）生成最终回答
    """
    # 构造系统提示词
    tool_result = state.get("tool_result", "")

    if tool_result:
        system_prompt = SystemMessage(
            content=f"""你是一个智能助手，基于工具返回的结果生成回答。

工具结果：{tool_result}

要求：
1. 必须使用工具结果中的信息，不要编造
2. 用自然、友好的语言回复用户
3. 保持简洁明了"""
        )
    else:
        system_prompt = SystemMessage(
            content="你是一个智能助手，直接基于常识和知识回答用户问题，保持友好和专业。"
        )

    # 组合消息
    messages = [system_prompt] + state["messages"]

    # 调用 LLM 生成回答
    ai_response = llm.invoke(messages)

    return {"messages": [ai_response]}


# ========== 5. 定义路由函数 ==========
def tool_router(state: ToolAgentState) -> Literal["weather_tool", "calculator_tool", "generate_answer"]:
    """
    条件路由：根据 tool_needed 决定跳转到哪个节点
    """
    tool_needed = state.get("tool_needed", "none")

    if tool_needed == "weather":
        return "weather_tool"
    elif tool_needed == "calculator":
        return "calculator_tool"
    else:
        return "generate_answer"


# ========== 6. 构建 Graph ==========
graph_builder = StateGraph(ToolAgentState)

# 添加所有节点
graph_builder.add_node("intent_recognition", intent_recognition_node)  # 意图识别
graph_builder.add_node("weather_tool", weather_tool_node)  # 天气工具
graph_builder.add_node("calculator_tool", calculator_tool_node)  # 计算器工具
graph_builder.add_node("generate_answer", generate_answer_node)  # 生成回答

# 添加边
graph_builder.add_edge(START, "intent_recognition")  # 开始 → 意图识别

# 条件边：从意图识别根据结果路由到不同工具或直接回答
graph_builder.add_conditional_edges(
    source="intent_recognition",
    path=tool_router,
    path_map={
        "weather_tool": "weather_tool",
        "calculator_tool": "calculator_tool",
        "generate_answer": "generate_answer"
    }
)

# 工具节点都指向回答生成节点
graph_builder.add_edge("weather_tool", "generate_answer")
graph_builder.add_edge("calculator_tool", "generate_answer")

# 回答生成节点指向结束
graph_builder.add_edge("generate_answer", END)

# 编译图
tool_agent_graph = graph_builder.compile()

# ========== 7. 可视化图结构 ==========
try:
    from IPython.display import Image, display

    png_data = tool_agent_graph.get_graph().draw_mermaid_png()
    with open("graph-tool-agent.png", "wb") as f:
        f.write(png_data)
    print("✅ 流程图已保存为 graph-tool-agent.png\n")

    # 如果在 Jupyter 环境中显示
    try:
        display(Image(data=png_data))
    except:
        pass
except Exception as e:
    print(f"⚠️  可视化失败（非必要）：{e}\n")


# ========== 8. 交互逻辑 ==========
def run_tool_agent():
    """启动多工具智能体"""
    print("=" * 60)
    print("🤖 多工具智能体已启动（天气查询 + 计算器）")
    print("=" * 60)
    print("功能说明：")
    print("  📍 天气查询：支持北京、上海、深圳、广州、杭州")
    print("  🔢 计算器：支持基本四则运算（+、-、*、/、括号）")
    print("  💬 其他问题：智能对话")
    print("\n提示：输入空字符串或 'quit' 退出\n")

    while True:
        # 获取用户输入
        user_input = input("👤 User: ")

        if user_input.strip() == "" or user_input.lower() == "quit":
            print("\n👋 对话结束，再见！")
            break

        # 初始化状态
        initial_state = {
            "messages": [{"role": "user", "content": user_input}]
        }

        # 执行图并流式输出
        print("🤖 Assistant: ", end="", flush=True)

        for event in tool_agent_graph.stream(initial_state):
            for node_output in event.values():
                # 提取并输出 AI 回答
                if "messages" in node_output:
                    latest_msg = node_output["messages"][-1]
                    if isinstance(latest_msg, AIMessage):
                        print(latest_msg.content)

        print()  # 换行


# ========== 9. 启动智能体 ==========
if __name__ == "__main__":
    run_tool_agent()
